from matplotlib import image
import torch
import torchvision
import torch.nn.functional as F
import torchvision.transforms as transforms
import torch.optim as optim
import matplotlib.pyplot as plt
import numpy as np
from PIL import Image
from torchvision import transforms

class  myNetwork(torch.nn.Module): # 定义自己的网络对象
    def  __init__(self) : 
        super(myNetwork, self).__init__() 
        self.convl1 = torch.nn.Conv2d(in_channels = 3, out_channels = 32, kernel_size = 5, stride = 1, padding = 2)
        self.pool11 = torch.nn.MaxPool2d(kernel_size = 2, stride = 2)
        self.convl2 = torch.nn.Conv2d(in_channels = 32, out_channels = 32, kernel_size = 5, stride = 1, padding = 2)
        self.pool12 = torch.nn.MaxPool2d(kernel_size = 2, stride = 2)
       
        self.convl3 = torch.nn.Conv2d(in_channels = 32, out_channels = 32, kernel_size = 5, stride = 1, padding = 2)
        self.pool13 = torch.nn.MaxPool2d(kernel_size = 2, stride = 2)
        self.fc1 = torch.nn.Linear(in_features = 32*5*5, out_features = 64, bias=True)
        self.fc2 = torch.nn.Linear(in_features = 64, out_features = 64, bias=True)
        self.fc3 = torch.nn.Linear(in_features = 64, out_features = 10, bias=True)
    def forward(self,train_data):
        output = self.pool11(torch.nn.functional.relu(self.convl1(train_data)))
        output = self.pool12(torch.nn.functional.relu(self.convl2(output)))
        output = self.pool13(torch.nn.functional.relu(self.convl3(output)))
        output = output.view(-1, 32*5*5)# 输出可视化
        output = torch.nn.functional.relu(self.fc1(output))
        output = torch.nn.functional.relu(self.fc2(output))
        output = self.fc3(output)
        return output


myTransforms  = transforms.Compose(
    [transforms.Resize((40,40)),
    transforms.ToTensor(),  # 将图像归一化并且变成张量
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]) 

predicted = 1
def prediect(img_path):
    net=torch.load('model.pkl')
    torch.no_grad()
    img=Image.open(img_path)
    img=myTransforms(img).unsqueeze(0)
    outputs = net(img)
    _, predicted = torch.max(outputs, 1)
if __name__ == '__main__':
    with open("result_lyx","w") as file:
        for i in range(20000):
            net=torch.load('model2.pkl')
            torch.no_grad()
            img=Image.open('/home/lyx/Desktop/LYX_python/dataset/tldataset/test/' + str(i) + ".png")
            img=myTransforms(img).unsqueeze(0)
            outputs = net(img)
            _, predicted = torch.max(outputs, 1)
            #print(str(i) + ".png" + " " + str(predicted)  +type[predicted] + "\n")
            file.write(str(i) + ".png" + " " + str(predicted.item()))
            file.write("\n")